package com.force.langchain4j;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.date.StopWatch;
import cn.hutool.core.io.resource.ClassPathResource;
import com.force.langchain4j.filesource.FileBaseLoader;
import com.force.langchain4j.vstore.impl.VectorStoreServiceImpl;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
//import io.milvus.client.MilvusServiceClient;
//import io.milvus.param.ConnectParam;
//import org.springframework.ai.document.Document;
//import org.springframework.ai.transformer.splitter.TokenTextSplitter;

import java.util.ArrayList;
import java.util.List;

public class JlamaMilvusExample {

    public static void main(String[] args) throws InterruptedException {
//        EmbeddingModel embeddingModel = VectorStoreServiceImpl.getEmbeddingModel("m3","sk-iwovlffhziuprtlelewtfsbxvligzuxbolncgfcphihjnyeo","https://api.siliconflow.cn/v1/");


//        MilvusServiceClient customMilvusClient = new MilvusServiceClient(
//                ConnectParam.newBuilder()
//                        .withHost("47.98.126.125")
//                        .withPort(19530)
//                        .withAuthorization("vdb", "vdb123")
//                        .build());

//        MilvusEmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
//                .milvusClient(customMilvusClient)
//                .collectionName("douluo_dalu")      // Name of the collection
//                .dimension(1024)                            // Dimension of vectors
//                .indexType(IndexType.FLAT)                 // Index type
//                .metricType(MetricType.COSINE)             // Metric type
//                .consistencyLevel(ConsistencyLevelEnum.EVENTUALLY)  // Consistency level
//                .autoFlushOnInsert(true)                   // Auto flush after insert
//                .idFieldName("id")                         // ID field name
//                .textFieldName("text")                     // Text field name
//                .metadataFieldName("metadata")             // Metadata field name
//                .vectorFieldName("vector")                 // Vector field name
//                .build();                                  // Build the MilvusEmbeddingStore instance
//        String chromaEndpoint = "http://localhost:8000";
//        ChromaEmbeddingStore embeddingStore = ChromaEmbeddingStore
//                .builder()
//                .baseUrl(chromaEndpoint)
//                .collectionName("douluo_dalu")
//                .logRequests(true)
//                .logResponses(true)
//                .build();
//
        EmbeddingModel miniLmL6V2EmbeddingModel = new AllMiniLmL6V2EmbeddingModel();

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .prefix("douluo_doc:")
                .dimension(384)
                .build();


//        ClassPathResource classPathResource = new ClassPathResource("斗罗大陆.txt");
//        Object load = FileBaseLoader.load(classPathResource.getAbsolutePath());
//        // 文档分割
//        TokenTextSplitter splitter = new TokenTextSplitter();
//        List<Document> splitDocuments = splitter.apply((List<Document>) load);
//        System.out.println("文件数量："+splitDocuments.size());
//        StopWatch stopWatch=new StopWatch();
//        stopWatch.start("开始存储到redis");
//        splitDocuments.parallelStream().forEach(item->{
//                TextSegment from = TextSegment.from(item.getText());
//                Embedding embedding = miniLmL6V2EmbeddingModel.embed(from).content();
//                embeddingStore.add(embedding, from);
//        });
//        stopWatch.prettyPrint();

//        TextSegment segment1 = TextSegment.from("I like football.");
//        Embedding embedding1 = embeddingModel.embed(segment1).content();
//        embeddingStore.add(embedding1, segment1);

//        TimeUnit.SECONDS.sleep(60);

//        TextSegment segment2 = TextSegment.from("The weather is good today.");
//        Embedding embedding2 = embeddingModel.embed(segment2).content();
//
//        embeddingStore.add(embedding2, segment2);

//        TimeUnit.SECONDS.sleep(60);

        String userQuery = "吞噬星空";
        Embedding queryEmbedding = miniLmL6V2EmbeddingModel.embed(userQuery).content();
        int maxResults = 1;

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(maxResults)
                .build();

        List<EmbeddingMatch<TextSegment>> matches = embeddingStore.search(embeddingSearchRequest).matches();
        List<String> results = new ArrayList<>();
        matches.forEach(embeddingMatch -> results.add(embeddingMatch.embedded().text()));


        System.out.println("Question: " + userQuery); // What is your favourite sport?
        System.out.println("Response: " + String.join("",results)); // I like football.
    }



}
